package tracing

import (
	"github.com/go-kratos/kratos/v2/middleware"
	"github.com/go-kratos/kratos/v2/middleware/tracing"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/exporters/jaeger"
	"go.opentelemetry.io/otel/sdk/resource"
	tracesdk "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
	"go.opentelemetry.io/otel/trace"

	"gitee.com/shuokeyun/kratos/conf"
)

// Server 链路跟踪中间件
func Server(c *conf.Trace, opts ...tracing.Option) middleware.Middleware {
	tp, err := tracerProvider(c)
	if err != nil {
		panic(err)
	}
	otel.SetTracerProvider(tp)
	return tracing.Server(opts...)
}

func tracerProvider(c *conf.Trace) (trace.TracerProvider, error) {
	// Create the Jaeger exporter
	exp, err := jaeger.New(jaeger.WithCollectorEndpoint(jaeger.WithEndpoint(c.GetJaeger().GetEndpoint())))
	if err != nil {
		return nil, err
	}

	kv := make([]attribute.KeyValue, 0, len(c.GetJaeger().GetAttribute())+1)
	kv = append(kv, semconv.ServiceNameKey.String(c.GetJaeger().GetServiceName()))
	for k, v := range c.GetJaeger().GetAttribute() {
		kv = append(kv, attribute.String(k, v))
	}
	ratio := c.GetJaeger().GetRatio()
	if ratio <= 0.0 {
		ratio = 1.0
	}
	tp := tracesdk.NewTracerProvider(
		tracesdk.WithSampler(tracesdk.ParentBased(tracesdk.TraceIDRatioBased(ratio))),
		// Always be sure to batch in production.
		tracesdk.WithBatcher(exp),
		tracesdk.WithResource(resource.NewWithAttributes(
			semconv.SchemaURL,
			kv...,
		)),
	)
	return tp, nil
}
